package com.epoint.flinkdemo.state;

import com.epoint.flinkdemo.entity.SenSorReading;
import com.epoint.flinkdemo.map.SensorReadingMapFunction;
import com.epoint.flinkdemo.source.WordCountFlat;
import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.state.BroadcastState;
import org.apache.flink.api.common.state.MapStateDescriptor;
import org.apache.flink.api.common.state.ReadOnlyBroadcastState;
import org.apache.flink.api.common.typeinfo.BasicTypeInfo;
import org.apache.flink.api.java.typeutils.MapTypeInfo;
import org.apache.flink.streaming.api.datastream.BroadcastStream;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.co.BroadcastProcessFunction;
import org.apache.flink.util.Collector;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * 测试广播状态
 */
public class BroadcastStateTest {

    private static final Logger LOG = LoggerFactory.getLogger(BroadcastStateTest.class);

    public static void main(String[] args) throws Exception{

        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        DataStreamSource<String> configSource = env.socketTextStream("localhost", 7788);
        MapStateDescriptor<String, Double> descriptor = new MapStateDescriptor<String, Double>("broadCastConfig",
                String.class, Double.class);
        BroadcastStream<Double> broadcastConfig = configSource.map(new MapFunction<String, Double>() {
            @Override
            public Double map(String s) throws Exception {
                return Double.parseDouble(s);
            }
        }).broadcast(descriptor);

        DataStreamSource<String> elementSource = env.socketTextStream("localhost", 7777);
        SingleOutputStreamOperator<SenSorReading> mapSource = elementSource.map(new SensorReadingMapFunction());

        // 必须是数据流connect广播流
        DataStream<SenSorReading> result = mapSource.connect(broadcastConfig).process(new BroadcastProcessFunction<SenSorReading, Double, SenSorReading>() {

            private Double defaultAlarm = 40.0;

            @Override
            public void processElement(SenSorReading value, ReadOnlyContext ctx, Collector<SenSorReading> out) throws Exception {
                // 处理数据流的数据
                // 从广播状态中获取数据
                ReadOnlyBroadcastState<String, Double> broadcastState = ctx.getBroadcastState(descriptor);
                if (broadcastState.contains("alarm")) {
                    defaultAlarm = broadcastState.get("alarm").doubleValue();
                }
                if (value.getTemp() > defaultAlarm) {
                    LOG.error("sensoreading temp {} > {}", value.getTemp(), defaultAlarm);
                    out.collect(value);
                }
            }

            @Override
            public void processBroadcastElement(Double value, Context ctx, Collector<SenSorReading> out) throws Exception {
                // 处理广播流的数据，将广播流数据保存至BroadcastState中
                BroadcastState<String, Double> broadcastState = ctx.getBroadcastState(descriptor);

                if (broadcastState.contains("alram")) {
                    LOG.info("old alarm is {}", broadcastState.get("alarm"));
                } else {
                    LOG.info("do not find old State, put first counterState {}", value);
                }
                broadcastState.put("alarm", value);
            }
        });
        result.print();
        env.execute("broadcast test job");

    }
}
